#-------------------------------------------------------------------------------
# Name:        dell_info.py
# Purpose:     Scan a network and retrieve serial numbers and hostnames
#
# Author:      robert
#
# Created:     02/08/2012
# Copyright:   (c) robert 2012
# Licence:     <your licence>
#-------------------------------------------------------------------------------

import sys
import re
import wmi
import csv
import string
import os
import socket


""" Establish class for WMI data. """

class wmiData(object):
    def __init__(self, host):
        print ' Connecting to WMI.'
        self.host = host
        self.target = wmi.WMI(host, user=uname, password=pword)

    def __str__(self):
        return 'WMI object, successfully connected to host {}.'.format(self.host.upper())

    def bios(self):
        """ Get BIOS info. """
        _bios = str(self.target.Win32_BIOS()[0])
        print '  BIOS data recovered.'
        return self.dataFormat(_bios)

    def freedisk_alert(self, host):
        """ Get free space and total space for each local drive. """
        _alert = [(host, disk.Caption, (long(disk.FreeSpace) * 1.0 / 1024**3) / (long(disk.Size) * 1.0 / 1024**3))
                  for disk in self.target.Win32_LogicalDisk(DriveType=3)
                  if (long(disk.FreeSpace) * 1.0 / 1024**3) / (long(disk.Size) * 1.0 / 1024**3) < 0.25]
        return _alert if _alert != [] else None

    def netdrives(self):
        """ Get drives shared from local machine. """
        return {share.Name:share.Path for share in self.target.Win32_Share()}

    def dataFormat(self, data):
        trans = re.sub(r'^[A-Za-z0-9_ \t\n]+{', '', data, 1).translate(None, '\t\n"').split(';')
        return zip(*[s.split(' = ') for s in trans if ' = ' in s])

    def proc(self):
        """ Get processor data. """
        _proc = str(self.target.Win32_Processor()[0])
        print '  Processor data recovered.'
        return self.dataFormat(_proc)

    """ This WMI data is specifically for a single user. """

    def processes(self):
        """ Get list of running processes. """
        category = ['Running Processes']
        proc_header = ['Name', 'PID', 'ExecutablePath', 'WorkingSetSize']
        proc_list = [(p.Name, p.ProcessId, p.ExecutablePath, p.WorkingSetSize)
                     for p in self.target.Win32_Process()]
        print '  Captured running processes.'
        return category, proc_header, proc_list

    def services(self):
        """ Get a list of services and their status. """
        category = ['Services']
        serv_header = ['Name', 'System Name', 'Status', 'Start Mode', 'Start User',
                       'Path', 'Accepts Pause?', 'Accepts Stop?']
        serv_list = [(s.Caption, s.Name, s.State, s.StartMode, s.StartName, s.PathName,
                     s.AcceptPause, s.AcceptStop) for s in self.target.Win32_Service()]
        print '  Service data recovered.'
        return category, serv_header, serv_list

    def startup(self):
        """ List startup applications. """
        category = ['Startup Applications']
        start_header = ['Location', 'Caption', 'Command']
        start_list = [(s.Location, s.Caption, s.Command)
                      for s in self.target.Win32_StartupCommand()]
        print '  Startup application data recovered.'
        return category, start_header, start_list

    def netinfo(self):
        """ Network interface information. """
        pass

    def eventlog(self):
        """ Event log errors and warnings:
            For individual user files, list all errors and warnings for last n days.
            For summary file, list number of errors. """
        print '  Capturing event logs.'
        _query = ('''Select Logfile, SourceName, TimeWritten, EventCode, Message
        from Win32_NTLogEvent where Type=\"error\" AND (Logfile=\"Application\" OR Logfile=\"System\") AND TimeWritten>=\"''' + fromdate + '"').encode('utf-8')
        errors = self.target.query(_query)
        print '   Captured {} logs.'.format(len(errors))
        if len(errors) == 0: return None, None
        e_list = [[e.split(' = ')[1] for e in entry if ' = ' in e]
                     for entry in [str(error).translate(None, '\t\n"').split(';')
                                   for error in errors]]
        e_header = [e.split(' = ')[0] for e in re.sub(r'^[\t\nA-Za-z0-9_ ]+{', '', str(errors[0]), 1).translate(None, '\t\n"').split(';') if ' = ' in e]
        return e_header, e_list


""" Establish class for host targets. """

class getNet(object):
    def __init__(self):
        self.hosts = []

    def hostfile(self, filename=None):
        """ Parse host file for targets. """
        if filename == None:
            filename = responses.ask('Enter filename:', 'file')
        print 'Parsing {} now.'.format(str(filename).upper())
        with open(filename, 'rb') as f:
            parse = csv.reader(f)
            for line in parse:
                try:
                    self.hosts.append(line[0])
                except IndexError:
                    break
        print 'Complete.'
        return self.hosts

    def netscan(self):
        """ Scan the network for targets using nmap.
            If multiple interfaces exist on the local machine you choose which to use. """
        f = open('acquired_hosts.csv', 'wb')
        csvout = csv.writer(f)
        host = socket.gethostname()
        ips = [ip for ip in socket.gethostbyname_ex(host)[2]
               if not ip.startswith('127.')]
        networks = {str(i + 1):'%s.%s.%s.0/24' % re.search(r'(\d+)\.(\d+)\.(\d+)\.\d+', ips[i]).groups()
                    for i in range(len(ips))}
        if len(networks) > 1:
            print 'The following networks were found:'
            print '\n'.join(['%s. %s' % (k,v) for k,v in networks.items()])
            choice = responses.ask('Please choose a network', 'multi', [str(i) for i in range(len(networks.items()))])
            if choice in networks.keys():
                netscan = networks[choice]
            elif choice == '':
                netscan = networks['1']
            else:
                print 'Choice out of range.'
                sys.exit()
        elif len(networks) == 1:
            netscan = networks['1']
            print '%s network found, setting as scan target.\n' % netscan
        else:
            print 'No networks found.  Quitting.'
            sys.exit()
        try:
            import nmap
            nm = nmap.PortScanner()
        except WindowsError:
            print 'Unable to initialize nmap.  Please check installation.'
            sys.exit()
        nscan = nm.scan(hosts=netscan, arguments='-sS -p 135')
        self.hosts = [nscan['scan'][s]['hostname'] for s in nscan['scan']
                      if 'open' in nscan['scan'][s]['tcp'][135]['state']]
        for host in self.hosts:
            csvout.writerow(host)
        f.close()
        return self.hosts

    def build_hosts(self):
        """ Create a target file by entering ip/hostname pairs. """
        saveit = responses.ask('Save to disk? (y/n)', 'binary')
        while True:
            hname = responses.ask('Enter host name, leave blank and press Enter if done.', 'text')
            if hname == '':
                break
            else:
                self.hosts.append(hname)
        print ''
        print self.hosts
        print 'Host list built.'
        if saveit == True:
            with open('build_hosts.csv', 'wb') as f:
                csv.writer(f).writerows(self.hosts)
        return self.hosts

    def get_ip(self, host):
        return socket.gethostbyname(host)


""" Class to take questions and return answers.  Record answers. """

class GetInput(object):
    def __init__(self):
        self.qdict = {}
        self.type_dict = {'binary': self._binary,
                          'file': self._fname,
                          'multi': self._multi,
                          'text': self._text
                          }

    def add(self, question, answer):
        self.qdict[question] = answer

    def ask(self, question, atype, akey=None):
        """ Takes a question to ask and response type.  Returns an answer. """
        self.question = question
        self.atype = atype
        self.akey = akey
        try:
            print ''
            print question
            answer = raw_input('>')
        except KeyboardInterrupt:
            print 'Input not recognized. Input valid response or type \'q\' to quit.'
            self.ask(self.question, self.atype, self.akey)
        else:
            if answer == 'q': sys.exit()
            if self.akey and answer not in self.akey:
                print 'Response not in accepted answer list.'
                self.ask(self.question, self.type, self.akey)
            self.qdict[question] = answer
            return (self.type_dict[atype](answer, akey))

    def _multi(self, answer, akey):
        """ Ask the user a multiple choice question. """
        if answer not in akey:
            print 'Response is not acceptable.'
            print ''
            return self.ask(self.question, 'multi', akey)
        return answer


    def _binary(self, answer, akey=None):
        """ Ask the user a question with a yes or no response.
            Store the question and answer in self.qdict. """
        pos = ['Y', 'y', 'YES', 'yes', 'Yes']
        neg = ['N', 'n', 'NO', 'No', 'no']
        if answer in pos: return True
        elif answer in neg: return False
        else:
            print 'Response not in accepted answer list.'
            print ''
            return self.ask(self.question, 'binary')

    def _fname(self, answer, akey=None):
        """ Ask the user a question where the answer is a filename.
            Verify the filename exists.
            Store the question and answer in self.qdict. """
        if '.csv' not in answer:
            print 'Only CSV files are accepted.'
            print ''
            return self.ask(self.question, 'file')
        elif not os.path.exists(answer):
            print 'File not found.'
            print ''
            return self.ask(self.question, 'file')
        else: return answer

    def _text(self, answer, akey=None):
        """ Return raw text. """
        return answer

""" Class to parse configuration files. """

class Configure(object):
    def __init__(self):
        self.config= {'FullScan': ['not configured', 'binary'],
                      'EventLog': ['not configured', 'binary'],
                      'DoServers': ['not configured', 'binary'],
                      'HostFile': ['not configured', 'file'],
                      'RefreshSummary': ['not configured', 'binary']}

    def check_file(self):
        if responses.ask('Load configuration file?', 'binary'):
            self.read()
            self.parse()
        else:
            self.do_config()

    def read(self):
        try:
            print os.getcwd()
            with open('scan_config.txt', 'rU') as f:
                self.cfile = f.readlines()
        except:
            option = responses.ask('Unable to open config file - (r)efresh, (q)uit or (c)ontinue?', 'multi',
                                   {'refresh': 'refresh', 'r': 'refresh', 'quit': 'quit', 'q': 'quit', 'continue': 'continue', 'c': 'continue'})
            if option == 'quit': sys.exit()
            return self.read() if option == 'refresh' else None
        else:
            print 'File read successfully.'

    def do_config(self):
        for option in self.config:
            if isinstance(self.config[option], list):
                self.config[option] = responses.ask(option, self.config[option][1])
            else:
                continue
        print 'Configuration complete:\n', '\n'.join([str(item) for item in self.config.items()])

    def write(self):
        pass

    def parse(self):
        _config = {line.split('=')[0].strip(): line.split('=')[1].strip() for line in self.cfile}
        for element in _config:
            if _config[element] == 'True': _config[element] = True
            elif _config[element] == 'False': _config[element] = False
        print 'Configuration found:\n'
        print '\n'.join([str(item) for item in _config.items()])
        load = responses.ask('Load this configuration? (y/n)> ', 'binary')
        if load: self.config.update(_config)
        self.do_config()



""" The following functions do work and retrieve attributes from the wmiData class. """

def dell_info(filename='scan_summary.csv'):
    """ Retrieve requested data. """
    hosts, oldscan = prep_hosts()
    print 'Found {} hosts to scan.'.format(len(hosts))
    while True:
        try:
            csvout = open(filename, 'wb')
        except:
            print "There's a problem with the output file.\nPlease close the output file if open, or type QUIT to quit."
            t = raw_input('>')
            if t in ['QUIT', 'quit', 'Quit', 'q', 'Q']: sys.exit()
        else:
            break
    outgoing = csv.writer(csvout)
    header = ['Hostname', 'Date']
    badhosts = []
    disk_alert = []
    do_header = True
    raw_input('Press Enter to continue or escape to break out.')
    for host in hosts:
        print 'Doing host {}'.format(host)
        row = [host, scandate]
        try:
            target = wmiData(host)
        except:
            print ' *** Connection failed. ***'
            if config.config['RefreshSummary']:
                print ' *** Writing existing date, if possible. ***'
                if oldscan[0][0] == host:
                    outgoing.writerow(oldscan.pop(0))
                else:
                    continue
            badhosts.append(host)
            continue
        new_header, data = each_row(target)
        if config.config['FullScan']:
            full_user(target, host)
        row.extend(data)
        if do_header:
            header.extend(new_header)
            outgoing.writerow(header)
            do_header = False
        print '  Checking disk space.'
        if target.freedisk_alert(host): disk_alert.extend(target.freedisk_alert(host))
        outgoing.writerow(row)
    if badhosts:
        print ''
        print 'Could not connect to the following hosts:\n'
        print '\n'.join(badhosts)
        for _ in range(3): outgoing.writerow('')
        outgoing.writerow(['Could not connect to the following hosts'])
        for host in badhosts: outgoing.writerow(host)
    if disk_alert:
        print ''
        print 'The following machines are low on disk space:\n'
        print '\n'.join([name + 'Drive ' + drive + ' has ' + str(space) + ' pct free space.'
                         for name, drive, space in disk_alert])
        for _ in range(3): outgoing.writerow('')
        outgoing.writerow(['The following machines are low on disk space'])
        outgoing.writerow(['IP', 'Host', 'Drive', '% Free Space'])
        for disk in disk_alert: outgoing.writerow(disk)
    csvout.close()

def each_row(target):
    """ Build summary data for each target and return as a list that will be written. """
    opts = [target.bios, target.proc]
    header = []
    data = []
    for opt in opts:
        attributes, values = opt()
        header.extend(attributes)
        data.extend(values)
    return header, data

def full_user(target, host):
    """ For deeper data scan, create folders and files. """
    transtable = string.maketrans('/:', "-'")
    if not os.path.exists(host): os.mkdir(host)
    os.chdir(host)
    filename = host + ' - ' + scandate.translate(transtable) + '.csv'
    f = open(filename, 'wb')
    csvout = csv.writer(f)
    opts = [target.processes, target.services, target.startup]
    csvout.writerow([host, scandate])
    for opt in opts:
        category, header, data = opt()
        csvout.writerow(category)
        csvout.writerow(header)
        for datum in sorted(data):
            try:
                csvout.writerow(datum)
            except:
                print datum
                sys.exit()
        for _ in range(3): csvout.writerow('')
    f.close()
    if config.config['EventLog']:
        while True:
            try:
                f = open('event_errors.csv', 'wb')
            except:
                print 'Close the file.'
                raw_input('')
            else:
                break
        eventout = csv.writer(f)
        e_header, e_list = target.eventlog()
        if e_header != None:
            eventout.writerow(e_header)
            eventout.writerows(e_list)
            f.close()
    os.chdir('..')

def prep_hosts():
    """ Check if scan_summary.csv exists.  If it does, ask to update instead of a new scan. """
    if config.config['RefreshSummary']:
        if os.path.exists('scan_summary.csv'):
            host_targets = getNet()
            hosts = host_targets.hostfile('scan_summary.csv')[1:]
            print 'Hosts from file:\n', '\n'.join([str(host) for host in hosts])
            print ''
            with open('scan_summary.csv', 'rb') as f:
                oldscan = [line for line in csv.reader(f)][1:]
            if responses.ask('Load more hosts from file?', 'binary'):
                hosts_fromfile = host_targets.hostfile()
                print 'Hosts from file:\n', '\n'.join([str(host) for host in hosts_fromfile])
                hosts.extend([host for host in hosts_fromfile if host[1] not in [h[1] for h in hosts]])
            return hosts, oldscan
    hosts = get_hosts()
    return hosts, None

def get_hosts():
    """ If no previous summary exists, or if additional hosts are desired, gather hosts. """
    c = responses.ask('''Please choose from one of the following options:\n
    \t1. Load targets from CSV.
    \t2. Scan network with NMAP.
    \t3. Enter single target name or IP.
    \t4. Build hosts list.''', 'multi', ['1', '2', '3', '4'])
    hclass = getNet()
    if c == '3':
        return [responses.ask('Enter IP address or hostname to scan.', 'text')]
    if c == '1': return hclass.hostfile()
    elif c == '2': return hclass.netscan()
    elif c == '4': return hclass.build_hosts()
    else:
        print "Input not recognized."
        print ""
        get_hosts()

def main():
    """ Establish globals for all questions and dates.
        Also check for the main Scan Data folder:
            If the folder exists, set the working directory to the folder.
            If the folder does not exist, create it and set the working directory. """
    import datetime
    import getpass
    global uname
    global pword
    global fromdate
    global scandate
    global responses
    global config
    uname = raw_input('Enter username: ')
    pword = getpass.win_getpass()
    responses = GetInput()
    config = Configure()
    config.check_file()
    d_search = re.compile(r'^(\d{4})-(\d{2})-(\d{2}) (\d{2}:\d{2}:\d{2})')
    today = re.search(d_search, str(datetime.datetime.now())).groups()
    targetdate = re.search(d_search, str(datetime.datetime.now() - datetime.timedelta(days=1))).groups()
    fromdate = '{1}/{2}/{0} 00:00:00'.format(targetdate[0], targetdate[1], targetdate[2])
    scandate = '{1}/{2}/{0} {3}'.format(today[0], today[1], today[2], today[3])
    print ''
    if not os.path.exists('Scan Data'):
        os.mkdir('Scan Data')
        os.chdir('Scan Data')
    else:
        os.chdir('Scan Data')
    dell_info()

if __name__ == '__main__':
    main()
